from termcolor import colored
from datetime import datetime
import json
from pydoc import pipepager
import sys

author = {
    "name":"gl4ssesbo1",
    "twitter":"https://twitter.com/gl4ssesbo1",
    "github":"https://github.com/gl4ssesbo1",
    "blog":"https://www.pepperclipp.com/"
}

needs_creds = True

variables = {
	"SERVICE": {
		"value": "ec2",
		"required": "true",
     		"description":"The service that will be used to run the module. It cannot be changed."
	},
	"AMIID": {
		"value": "ami-014d8dccd70fd2632",
		"required": "true",
        "description":"The AMI of the system. The default is ubuntu/images/hvm-ssd/ubuntu-bionic-18.04-amd64-server-20200908"
	},
	"SUBNETID": {
		"value": "",
		"required": "true",
        "description":"ID of the subnet it will be part of. The subnet needs to exist."
	},
	"SECURITYGROUP": {
		"value": "",
		"required": "true",
        "description":"ID of the security group it will be part of. It also needs to exist."
	},
	"INSTANCETYPE": {
		"value": "t2.micro",
		"required": "true",
        "description":"Instance type. Default is t2.micro"
	},
    "KEYNAME": {
        "value": "",
        "required": "true",
        "description":"Name of the SSH key. It can exist or not. If not exist, it will be createdm but you need permission to create it."
    },
    "NUMBEROFMACHINES": {
        "value": "1",
        "required": "true",
        "description":"How many machines do you want. All will have the same key and the same user data. If you specify more instances than Amazon EC2 can launch in the target Availability Zone, Amazon EC2 launches the largest possible number of instances above MinCount which is 1."
    },
	"USERDATA": {
		"value": "",
		"required": "true",
        "description":"User data encoded in base64."
    }
}

description = "You must provide policies in JSON format in IAM. However, for AWS CloudFormation templates formatted in YAML, you can provide the policy in JSON or YAML format. AWS CloudFormation always converts a YAML policy to JSON format before submitting it to IAM."

aws_command = "aws ec2 run-instances --image-id <AMI ID> --instance-type <instance type> --count <amount of machines> --subnet-id <subnet-id> --key-name <key-name> --security-group-ids <security-group-ids> --user-data <user data base64 string or file using file://pathoffile> --region <region> --profile <profile>"

colors = [
    "not-used",
    "red",
    "blue",
    "yellow",
    "green",
    "magenta",
    "cyan",
    "white"
]

output = ""

def list_dictionary(d, n_tab):
	global output
	if isinstance(d, list):
		n_tab += 1
		for i in d:
			if not isinstance(i, list) and not isinstance(i, dict):
				output += ("{}{}\n".format("\t" * n_tab, colored(i, colors[n_tab])))
			else:
				list_dictionary(i, n_tab)
	elif isinstance(d, dict):
		n_tab+=1
		for key, value in d.items():
			if not isinstance(value, dict) and not isinstance(value, list):
				output += ("{}{}: {}\n".format("\t"*n_tab, colored(key, colors[n_tab], attrs=['bold']) , colored(value, colors[n_tab+1])))
			else:
				output += ("{}{}:\n".format("\t"*n_tab, colored(key, colors[n_tab], attrs=['bold'])))
				list_dictionary(value, n_tab)

def exploit(profile, workspace):
    now = datetime.now()
    dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
    file = "{}_ec2_create_instance".format(dt_string)
    filename = "./workspaces/{}/{}".format(workspace, file)

    global output

    full_output = open(filename, 'a')
    keyname = variables['KEYNAME']['value']
    count = int(variables['NUMBEROFMACHINES']['value'])

    try:
        response = profile.describe_key_pairs(
            KeyNames=[
                keyname,
            ]
        )
        resp = response['KeyPairs']
        for json_data in resp:
            output += "{}: \n".format(
                colored("----------------------------------", "yellow", attrs=['bold'])
            )
            output += "{}: {}\n".format(
                colored("KeyName", "yellow", attrs=['bold']),
                colored(json_data['KeyName'], "white", attrs=['bold']),
            )
            output += "{}: \n".format(
                colored("----------------------------------", "yellow", attrs=['bold'])
            )
            output += "\t{}: {}\n".format(
                colored("KeyPairId", "blue", attrs=['bold']),
                colored(json_data['KeyPairId'], "yellow", attrs=['bold']),
            )
            output += "\t{}: {}\n".format(
                colored("KeyFingerprint", "blue", attrs=['bold']),
                colored(json_data['KeyFingerprint'], "yellow", attrs=['bold']),
            )
            output += "\t{}: {}\n".format(
                colored("KeyName", "blue", attrs=['bold']),
                colored(json_data['KeyName'], "yellow", attrs=['bold']),
            )
            output += "\t{}: {}\n".format(
                colored("Tags", "blue", attrs=['bold']),
                colored(json_data['Tags'], "yellow", attrs=['bold']),
            )
            output += "{}: \n".format(
                colored("----------------------------------", "yellow", attrs=['bold'])
            )
        json.dump(json_data, full_output, indent=4, default=str)
        print(colored("[*] Key '{}' exists. Creating the instance...".format(keyname), "yellow"))

    except:
        try:
            print(colored("[*] Key '{}' does not exist. Creating the key...".format(keyname), "yellow"))
            keyfile = "./ssh_keys/{}".format(keyname)
            outfile = open(keyfile, 'w')
            key_pair = profile.create_key_pair(
                KeyName=keyname
            )
            del key_pair['ResponseMetadata']
            json_data = key_pair
            outfile.write(key_pair['KeyMaterial'])
            outfile.close()

            output += "{}: \n".format(
                colored("----------------------------------", "yellow", attrs=['bold'])
            )
            output += "{}: {}\n".format(
                colored("KeyName", "yellow", attrs=['bold']),
                colored(json_data['KeyName'], "white", attrs=['bold']),
            )
            output += "{}: \n".format(
                colored("----------------------------------", "yellow", attrs=['bold'])
            )
            output += "\t{}: {}\n".format(
                colored("KeyPairId", "blue", attrs=['bold']),
                colored(json_data['KeyPairId'], "yellow", attrs=['bold']),
            )
            output += "\t{}: {}\n".format(
                colored("KeyFingerprint", "blue", attrs=['bold']),
                colored(json_data['KeyFingerprint'], "yellow", attrs=['bold']),
            )
            output += "\t{}: {}\n".format(
                colored("KeyName", "blue", attrs=['bold']),
                colored(json_data['KeyName'], "yellow", attrs=['bold']),
            )
            output += "\t{}: {}\n".format(
                colored("KeyMaterial", "blue", attrs=['bold']),
                colored("Key saved to file '{}'.".format(keyfile), "green"),
            )
            output += "{}: \n".format(
                colored("----------------------------------", "yellow", attrs=['bold'])
            )
            with open(filename, 'w') as outfile:
                json.dump(key_pair, full_output, indent=4, default=str)

            print(colored("[*] Key saved to file '{}'.".format(keyfile), "green"))
            print(colored("[*] Creating the instance...", "yellow"))

        except:
            e = sys.exc_info()
            print(colored("[*] {}".format(e), "red"))

    try:
        instancetype = variables['INSTANCETYPE']['value']
        amiid = variables['AMIID']['value']
        subnetid = variables['SUBNETID']['value']
        securitygroup = variables['SECURITYGROUP']['value']
        userdata = variables['USERDATA']['value']

        response = profile.run_instances(
            BlockDeviceMappings=[
                {
                    'DeviceName': '/dev/xvda',
                    'Ebs': {

                        'DeleteOnTermination': True,
                        'VolumeSize': 8,
                        'VolumeType': 'gp2'
                    },
                },
            ],
            ImageId=amiid,
            InstanceType=instancetype,
            MaxCount=count,
            MinCount=1,
            UserData=userdata,
            Monitoring={
                'Enabled': False
            },
            SecurityGroupIds=[
                securitygroup,
            ],
            SubnetId=subnetid
        )
        del response['ResponseMetadata']
        output += "{}: {}\n".format(
            colored("OwnerId", "red", attrs=['bold']),
            colored(response['OwnerId'], "blue", attrs=['bold']),
        )
        output += "{}: {}\n".format(
            colored("ReservationId", "red", attrs=['bold']),
            colored(response['ReservationId'], "blue", attrs=['bold']),
        )
        output += "{}:\n".format(
            colored("Groups", "red", attrs=['bold'])
        )
        output += "{}:\n".format(
            colored("Instances", "red", attrs=['bold'])
        )
        for group in response['Groups']:
            output += "\t{}: {}\n".format(
                colored("GroupName", "blue", attrs=['bold']),
                colored(group['GroupName'], "yellow"),
            )
            output += "\t{}: {}\n".format(
                colored("GroupId", "blue", attrs=['bold']),
                colored(group['GroupId'], "yellow"),
            )
        for json_data in response['Instances']:
            json_data['UserData'] = userdata
            title_name = 'InstanceId'
            n_tab = 1
            if isinstance(json_data, list):
                output += colored("\t---------------------------------\n", "yellow", attrs=['bold'])
                for data in json_data:
                    output += colored("\t{}: {}\n".format(title_name, json_data[title_name]), "yellow", attrs=['bold'])
                    output += colored("\t---------------------------------\n", "yellow", attrs=['bold'])
                    list_dictionary(data, n_tab)
                    output += colored("\t---------------------------------\n", "yellow", attrs=['bold'])
            else:
                output += colored("\t---------------------------------\n", "yellow", attrs=['bold'])
                output += colored("\t{}: {}\n".format(title_name, json_data[title_name]), "yellow", attrs=['bold'])
                output += colored("\t---------------------------------\n", "yellow", attrs=['bold'])
                list_dictionary(json_data, n_tab)
                output += colored("\t---------------------------------\n", "yellow", attrs=['bold'])

        json.dump(response, full_output, indent=4, default=str)
        print(colored("[*] Output written to file", "green"), colored("'{}'".format(filename), "blue"),
              colored(".", "green"))

        for inst in response['Instances']:
            print(colored("[*] Instance '{}' with sshkey '{}' successfully created.".format(inst['InstanceId'], keyname),"green"))

        pipepager(output,"less -R")

        output = ""
        full_output.close()
    except:
        e = sys.exc_info()
        print(colored("[*] {}".format(e), "red"))
